#pragma once

#include "bencode.hpp"

namespace bitgene {

template <typename Stream, template <typename> class Traits>
Stream & operator >>(Stream &stream, basic_bencode<Traits> &r)
{
	typename Stream::int_type ch = stream.get();
	if (!stream.good())
		return stream;
	switch (ch)
	{
		case 'i':
		{
			// read until 'e'
			std::stringstream ss;
			if (!stream.get(*ss.rdbuf(), 'e').good())
				return stream;
			// skip 'e'
			stream.get();
			// convert the string into an integer
			typename basic_bencode<Traits>::integer_type v;
			ss >> std::noskipws >> v;
			if (ss.fail() || !ss.eof())
			{
				stream.setstate(std::ios::failbit);
				return stream;
			}
			// assign
			r = v;
			break;
		}
		case 'l':
		{
			typename basic_bencode<Traits>::list_type v;
			for (; stream.peek() != 'e'; )
			{
				basic_bencode<Traits> sub;
				if (!(stream >> sub).good())
					return stream;
				v.push_back(sub);
			}
			// skip the last 'e'
			stream.get();
			// assign
			r = v;
			break;
		}
		case 'd':
		{
			typedef typename basic_bencode<Traits>::dictionary_type dictionary_type;
			typedef typename dictionary_type::key_type key_type;
			dictionary_type v;
			for (; stream.peek() != 'e';)
			{
				basic_bencode<Traits> key, mapped;
				stream >> key >> mapped;
				if (!stream.good())
					return stream;
				key_type *pkey = key.template pointer<key_type>();
				if (pkey == 0)
				{
					// check the key_type of dictionaries
					stream.setstate(std::ios::failbit);
					return stream;
				}
				//v.insert(dictionary_type::value_type(*pkey, mapped));
				v[*pkey] = mapped;
			}
			// skip the last 'e'
			stream.get();
			// assign
			r = v;
			break;
		}
		default:
		{
			// put back the first character
			stream.putback(ch);
			std::stringstream ss;
			if (!stream.get(*ss.rdbuf(), ':').good())
				return stream;
			// skip ':'
			stream.get();
			// convert the string into an integer
			std::streamsize size;
			ss >> std::noskipws >> size;
			if (ss.fail() || !ss.eof())
			{
				stream.setstate(std::ios::failbit);
				return stream;
			}
			// read the content of the string
			typename basic_bencode<Traits>::string_type v;
			v.resize(size);
			if (!stream.read(&v[0], size).good())
				return stream;
			// assign
			r = v;
			break;
		}
	}
	return stream;
}

template <typename Stream, template <typename> class Traits>
Stream & operator <<(Stream &stream, const basic_bencode<Traits> &r)
{
	typedef typename basic_bencode<Traits>::integer_type integer_type;
	typedef typename basic_bencode<Traits>::string_type string_type;
	typedef typename basic_bencode<Traits>::list_type list_type;
	typedef typename basic_bencode<Traits>::dictionary_type dictionary_type;

	if (const integer_type *v = r.template const_pointer<integer_type>())
		stream << 'i' << *v << 'e';
	else if (const string_type *v = r.template const_pointer<string_type>())
		stream << v->size() << ':' << *v;
	else if (const list_type *v = r.template const_pointer<list_type>())
	{
		stream << 'l';
		for (typename list_type::const_iterator i = v->begin(); i != v->end(); ++i)
			stream << *i;
		stream << 'e';
	}
	else if (const dictionary_type *v = r.template const_pointer<dictionary_type>())
	{
		stream << 'd';
		for (typename dictionary_type::const_iterator i = v->begin(); i != v->end(); ++i)
			stream << i->first.size() << ':' << i->first << i->second;
		stream << 'e';
	}
	else
		stream.setstate(std::ios::failbit);
	return stream;
}

} // namespace bitgene
